
# coding: utf-8
"""
通过实现Grad-CAM学习module中的forward_hook和backward_hook函数
"""
import cv2
import os
import numpy as np
from PIL import Image
import torch
import torchvision.transforms as transforms
from model import  LeNet5
import matplotlib.pyplot as plt
import io
import predict_DMQ_Axial

def img_transform(img_in, transform):
    """
    将img进行预处理，并转换成模型输入所需的形式—— B*C*H*W
    :param img_roi: np.array
    :return:
    """
    img = img_in.copy()
    img = Image.fromarray(np.uint8(img))
    img = transform(img)
    img = img.unsqueeze(0)  # C*H*W --> B*C*H*W
    return img


def img_preprocess(img_in):
    """
    读取图片，转为模型可读的形式
    :param img_in: ndarray, [H, W, C]
    :return: PIL.image
    """
    img = img_in.copy()
    img = cv2.resize(img, (28,28))
    img = img[:, :, ::-1]  # BGR --> RGB
    transform = transforms.Compose([
        transforms.ToTensor()
    ])
    img_input = img_transform(img, transform)
    return img_input


def backward_hook(module, grad_in, grad_out):
    grad_block.append(grad_out[0].detach())


def farward_hook(module, input, output):
    fmap_block.append(output)



def show_cam_on_image(img, mask, out_dir):
    heatmap = cv2.applyColorMap(np.uint8(255 * mask), cv2.COLORMAP_JET)
    heatmap = np.float32(heatmap) / 255
    cam = heatmap + np.float32(img)
    cam = cam / np.max(cam)

    path_cam_img = os.path.join(out_dir, "cam.jpg")
    path_raw_img = os.path.join(out_dir, "raw.jpg")
    if not os.path.exists(out_dir):
        os.makedirs(out_dir)
    cv2.imwrite(path_cam_img, np.uint8(255 * cam))
    cv2.imwrite(path_raw_img, np.uint8(255 * img))


def comp_class_vec(ouput_vec, index=None):
    """
    计算类向量
    :param ouput_vec: tensor
    :param index: int，指定类别
    :return: tensor
    """
    if not index:
        index = np.argmax(ouput_vec.cpu().data.numpy())
    else:
        index = np.array(index)
    index = index[np.newaxis, np.newaxis]
    index = torch.from_numpy(index)
    one_hot = torch.zeros(1, 2).scatter_(1, index, 1)
    one_hot.requires_grad = True
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    class_vec = torch.sum(one_hot.to(device) * ouput_vec)  # one_hot = 11.8605

    return class_vec


def gen_cam(feature_map, grads):
    """
    依据梯度和特征图，生成cam
    :param feature_map: np.array， in [C, H, W]
    :param grads: np.array， in [C, H, W]
    :return: np.array, [H, W]
    """
    cam = np.zeros(feature_map.shape[1:], dtype=np.float32)  # cam shape (H, W)

    weights = np.mean(grads, axis=(1, 2))  #

    for i, w in enumerate(weights):
        cam += w * feature_map[i, :, :]

    cam = np.maximum(cam, 0)
    cam = cv2.resize(cam, (28,28))
    cam -= np.min(cam)
    cam /= np.max(cam)

    return cam

fmap_block = list()
grad_block = list()
def get_cam_img(path_img):
    filename = os.path.basename(path_img)
    patient = filename.split('.')[0]

    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    # path_img = "./pridiction/MMQ/Axial/" + patient + '.png'
    path_net='./LeNet5.pth'
    
    img = cv2.imread(path_img, 1)  # H*W*C
    img_input = img_preprocess(img)
    net = LeNet5(num_classes=2).to(device)
    net.load_state_dict(torch.load(path_net,map_location="cpu"))
    net.conv2.register_forward_hook(farward_hook)
    net.conv2.register_backward_hook(backward_hook)
    output = net(img_input.to(device))

    net.zero_grad()
    class_loss = comp_class_vec(output)
    class_loss.backward()

    grads_val = grad_block[0].cpu().data.numpy().squeeze()
    fmap = fmap_block[0].cpu().data.numpy().squeeze()
    cam = gen_cam(fmap, grads_val)

    img_show = np.float32(cv2.resize(img, (28,28))) / 255
    heatmap = cv2.applyColorMap(np.uint8(255 * cam), cv2.COLORMAP_JET)
    heatmap = np.float32(heatmap) / 255
    cam = heatmap + np.float32(img_show)
    cam = cam / np.max(cam)
    
    img = Image.fromarray(cv2.cvtColor(np.uint8(255 * cam),cv2.COLOR_BGR2RGB))

    result = predict_DMQ_Axial.main(path_img)
    prediction = result['prediction']
    prob = result['probability']
    original_class = result['original']

    print_res = "patient:{}     original:{}     predict: {}     prob: {:.3}".format(patient,original_class,prediction,
                                                 prob)

    
    fig = plt.figure()
    plt.axis('off')
    plt.imshow(img)
    plt.title(print_res)
    
    canvas = fig.canvas
    buffer = io.BytesIO()
    canvas.print_png(buffer)
    data = buffer.getvalue()
    buffer.close()

    return data

def get_raw_img(patient):
    path_img = "./pridiction/MMQ/Axial/" + patient + '.png'
    img = cv2.imread(path_img, 1)
    img_show = np.float32(cv2.resize(img, (28,28))) / 255

    img = Image.fromarray(np.uint8(255 * img_show))
    fig = plt.figure()
    plt.axis('off')
    plt.imshow(img)
    
    canvas = fig.canvas
    buffer = io.BytesIO()
    canvas.print_png(buffer)
    data = buffer.getvalue()
    buffer.close()
    return data







